package db

import (
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

// SSHTunnel SSH隧道结构
type SSHTunnel struct {
	sshClient   *ssh.Client
	listener    net.Listener
	localAddr   string
	remoteAddr  string
	sshAddr     string
	config      *ssh.ClientConfig
	mu          sync.Mutex
	closed      bool
	connections map[net.Conn]bool
	connMu      sync.RWMutex
}

// SSHTunnelConfig SSH隧道配置
type SSHTunnelConfig struct {
	SSHHost        string
	SSHPort        int
	SSHUser        string
	SSHPassword    string
	SSHPrivateKey  string
	SSHPassphrase  string
	LocalPort      int    // 0表示自动分配
	RemoteHost     string // MySQL服务器地址
	RemotePort     int    // MySQL端口
	KeepAlive      bool
	TCPNoDelay     bool
	SendBufferSize int // KB
	RecvBufferSize int // KB
}

// NewSSHTunnel 创建新的SSH隧道
func NewSSHTunnel(config SSHTunnelConfig) (*SSHTunnel, error) {
	// 构建SSH配置
	sshConfig, err := buildSSHConfig(config)
	if err != nil {
		return nil, fmt.Errorf("构建SSH配置失败: %w", err)
	}

	sshAddr := fmt.Sprintf("%s:%d", config.SSHHost, config.SSHPort)
	remoteAddr := fmt.Sprintf("%s:%d", config.RemoteHost, config.RemotePort)

	tunnel := &SSHTunnel{
		sshAddr:     sshAddr,
		remoteAddr:  remoteAddr,
		config:      sshConfig,
		connections: make(map[net.Conn]bool),
	}

	return tunnel, nil
}

// buildSSHConfig 构建SSH客户端配置
func buildSSHConfig(config SSHTunnelConfig) (*ssh.ClientConfig, error) {
	var authMethods []ssh.AuthMethod

	// 优先使用私钥认证
	if config.SSHPrivateKey != "" {
		keyAuth, err := getKeyAuth(config.SSHPrivateKey, config.SSHPassphrase)
		if err != nil {
			return nil, fmt.Errorf("读取私钥失败: %w", err)
		}
		authMethods = append(authMethods, keyAuth)
	}

	// SSH Agent认证
	if agentAuth := getAgentAuth(); agentAuth != nil {
		authMethods = append(authMethods, agentAuth)
	}

	// 密码认证
	if config.SSHPassword != "" {
		authMethods = append(authMethods, ssh.Password(config.SSHPassword))
	}

	if len(authMethods) == 0 {
		return nil, fmt.Errorf("未提供任何SSH认证方式")
	}

	sshConfig := &ssh.ClientConfig{
		User:            config.SSHUser,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 生产环境应该验证主机密钥
		Timeout:         30 * time.Second,
	}

	// 启用keepalive
	if config.KeepAlive {
		sshConfig.ClientVersion = "SSH-2.0-NavidogSSH_1.0"
	}

	return sshConfig, nil
}

// getKeyAuth 获取私钥认证方法
func getKeyAuth(keyPath, passphrase string) (ssh.AuthMethod, error) {
	// 支持相对路径和绝对路径
	if !filepath.IsAbs(keyPath) {
		home, err := os.UserHomeDir()
		if err != nil {
			return nil, fmt.Errorf("获取用户主目录失败: %w", err)
		}
		keyPath = filepath.Join(home, ".ssh", keyPath)
	}

	key, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("读取私钥文件失败: %w", err)
	}

	var signer ssh.Signer
	if passphrase != "" {
		signer, err = ssh.ParsePrivateKeyWithPassphrase(key, []byte(passphrase))
	} else {
		signer, err = ssh.ParsePrivateKey(key)
	}

	if err != nil {
		return nil, fmt.Errorf("解析私钥失败: %w", err)
	}

	return ssh.PublicKeys(signer), nil
}

// getAgentAuth 获取SSH Agent认证方法
func getAgentAuth() ssh.AuthMethod {
	if sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		return ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers)
	}
	return nil
}

// Start 启动SSH隧道
func (t *SSHTunnel) Start() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed {
		return fmt.Errorf("隧道已关闭")
	}

	// 连接SSH服务器
	sshClient, err := ssh.Dial("tcp", t.sshAddr, t.config)
	if err != nil {
		return fmt.Errorf("连接SSH服务器失败: %w", err)
	}

	t.sshClient = sshClient

	// 启用SSH keepalive
	go t.keepAlive()

	// 创建本地监听器
	listener, err := net.Listen("tcp", "127.0.0.1:0") // 自动分配端口
	if err != nil {
		sshClient.Close()
		return fmt.Errorf("创建本地监听器失败: %w", err)
	}

	t.listener = listener
	t.localAddr = listener.Addr().String()

	// 启动隧道转发
	go t.handleConnections()

	return nil
}

// keepAlive 维护SSH连接
func (t *SSHTunnel) keepAlive() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			t.mu.Lock()
			if t.closed || t.sshClient == nil {
				t.mu.Unlock()
				return
			}
			// 发送keepalive请求
			_, _, err := t.sshClient.SendRequest("keepalive@openssh.com", true, nil)
			t.mu.Unlock()

			if err != nil {
				// keepalive失败，可能需要重连
				fmt.Printf("SSH keepalive失败: %v\n", err)
				return
			}
		}
	}
}

// handleConnections 处理连接转发
func (t *SSHTunnel) handleConnections() {
	for {
		t.mu.Lock()
		if t.closed || t.listener == nil {
			t.mu.Unlock()
			return
		}
		listener := t.listener
		t.mu.Unlock()

		localConn, err := listener.Accept()
		if err != nil {
			if !t.closed {
				fmt.Printf("接受本地连接失败: %v\n", err)
			}
			return
		}

		// 异步处理每个连接
		go t.handleConnection(localConn)
	}
}

// handleConnection 处理单个连接的转发
func (t *SSHTunnel) handleConnection(localConn net.Conn) {
	// 记录连接
	t.connMu.Lock()
	t.connections[localConn] = true
	t.connMu.Unlock()

	defer func() {
		localConn.Close()
		t.connMu.Lock()
		delete(t.connections, localConn)
		t.connMu.Unlock()
	}()

	// 优化本地连接
	if err := optimizeConnection(localConn); err != nil {
		fmt.Printf("优化本地连接失败: %v\n", err)
	}

	// 通过SSH连接到远程服务器
	t.mu.Lock()
	if t.closed || t.sshClient == nil {
		t.mu.Unlock()
		return
	}
	sshClient := t.sshClient
	t.mu.Unlock()

	remoteConn, err := sshClient.Dial("tcp", t.remoteAddr)
	if err != nil {
		fmt.Printf("连接远程服务器失败: %v\n", err)
		return
	}
	defer remoteConn.Close()

	// 优化远程连接
	if err := optimizeConnection(remoteConn); err != nil {
		fmt.Printf("优化远程连接失败: %v\n", err)
	}

	// 使用零拷贝进行双向数据转发
	var wg sync.WaitGroup
	wg.Add(2)

	// 本地 -> 远程
	go func() {
		defer wg.Done()
		if err := copyWithZeroCopy(remoteConn, localConn); err != nil && err != io.EOF {
			// fmt.Printf("本地到远程转发失败: %v\n", err)
		}
		remoteConn.Close()
	}()

	// 远程 -> 本地
	go func() {
		defer wg.Done()
		if err := copyWithZeroCopy(localConn, remoteConn); err != nil && err != io.EOF {
			// fmt.Printf("远程到本地转发失败: %v\n", err)
		}
		localConn.Close()
	}()

	wg.Wait()
}

// optimizeConnection 优化网络连接
func optimizeConnection(conn net.Conn) error {
	if tcpConn, ok := conn.(*net.TCPConn); ok {
		// 启用TCP_NODELAY，禁用Nagle算法
		if err := tcpConn.SetNoDelay(true); err != nil {
			return fmt.Errorf("设置TCP_NODELAY失败: %w", err)
		}

		// 设置keepalive
		if err := tcpConn.SetKeepAlive(true); err != nil {
			return fmt.Errorf("设置keepalive失败: %w", err)
		}

		// 设置keepalive参数
		if err := tcpConn.SetKeepAlivePeriod(30 * time.Second); err != nil {
			return fmt.Errorf("设置keepalive周期失败: %w", err)
		}

		// 设置读写缓冲区
		if err := tcpConn.SetReadBuffer(64 * 1024); err != nil {
			return fmt.Errorf("设置读缓冲区失败: %w", err)
		}

		if err := tcpConn.SetWriteBuffer(64 * 1024); err != nil {
			return fmt.Errorf("设置写缓冲区失败: %w", err)
		}
	}
	return nil
}

// copyWithZeroCopy 使用零拷贝技术转发数据
func copyWithZeroCopy(dst io.Writer, src io.Reader) error {
	// 尝试使用系统调用进行零拷贝
	if srcFile, srcOk := src.(interface{ File() (*os.File, error) }); srcOk {
		if dstFile, dstOk := dst.(interface{ File() (*os.File, error) }); dstOk {
			if srcF, err := srcFile.File(); err == nil {
				if dstF, err := dstFile.File(); err == nil {
					defer srcF.Close()
					defer dstF.Close()
					// 在Linux上可以使用splice系统调用实现零拷贝
					// 这里使用标准库的CopyBuffer作为fallback
					return copyWithBuffer(dst, src)
				}
			}
		}
	}

	// Fallback到缓冲区拷贝
	return copyWithBuffer(dst, src)
}

// copyWithBuffer 使用大缓冲区进行数据拷贝
func copyWithBuffer(dst io.Writer, src io.Reader) error {
	// 使用64KB缓冲区以提高性能
	buf := make([]byte, 64*1024)
	_, err := io.CopyBuffer(dst, src, buf)
	return err
}

// GetLocalAddr 获取本地监听地址
func (t *SSHTunnel) GetLocalAddr() string {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.localAddr
}

// Close 关闭SSH隧道
func (t *SSHTunnel) Close() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed {
		return nil
	}

	t.closed = true

	// 关闭所有活跃连接
	t.connMu.Lock()
	for conn := range t.connections {
		conn.Close()
	}
	t.connMu.Unlock()

	// 关闭监听器
	if t.listener != nil {
		t.listener.Close()
	}

	// 关闭SSH客户端
	if t.sshClient != nil {
		t.sshClient.Close()
	}

	return nil
}

// IsAlive 检查隧道是否存活
func (t *SSHTunnel) IsAlive() bool {
	t.mu.Lock()
	defer t.mu.Unlock()

	if t.closed || t.sshClient == nil {
		return false
	}

	// 简单的连接检查
	_, _, err := t.sshClient.SendRequest("keepalive@openssh.com", true, nil)
	return err == nil
}
